"""
估值图表模块 - 提供估值指标图表生成功能
"""

import pandas as pd
import plotly.graph_objects as go
from typing import List, Dict, Any, Optional

def create_valuation_chart(
    df: pd.DataFrame, 
    title: str = "估值指标", 
    height: int = 500,
    pe_col: str = 'PE(TTM)',
    pb_col: str = '市净率',
    date_col: str = '数据日期'
) -> go.Figure:
    """
    创建估值指标图表
    
    参数:
        df: 包含估值数据的DataFrame，必须包含 date_col, pe_col, pb_col 列
        title: 图表标题
        height: 图表高度
        pe_col: PE列名
        pb_col: PB列名
        date_col: 日期列名
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = [date_col, pe_col, pb_col]
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 创建图表
    fig = go.Figure()
    
    # 添加PE线
    fig.add_trace(go.Scatter(
        x=df[date_col],
        y=df[pe_col],
        mode='lines+markers',
        name=pe_col,
        line=dict(color='blue')
    ))
    
    # 添加PB线
    fig.add_trace(go.Scatter(
        x=df[date_col],
        y=df[pb_col],
        mode='lines+markers',
        name=pb_col,
        line=dict(color='red'),
        yaxis="y2"
    ))
    
    # 设置图表布局
    fig.update_layout(
        title=title,
        xaxis_title="日期",
        yaxis_title=pe_col,
        yaxis2=dict(
            title=pb_col,
            overlaying='y',
            side='right'
        ),
        height=height,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    return fig

def create_valuation_comparison_chart(
    df: pd.DataFrame, 
    company_col: str,
    pe_col: str = 'PE(TTM)',
    pb_col: str = '市净率',
    title: str = "估值对比", 
    height: int = 500
) -> go.Figure:
    """
    创建估值对比图表
    
    参数:
        df: 包含多家公司估值数据的DataFrame，必须包含 company_col, pe_col, pb_col 列
        company_col: 公司名称列
        pe_col: PE列名
        pb_col: PB列名
        title: 图表标题
        height: 图表高度
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = [company_col, pe_col, pb_col]
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 创建图表
    fig = go.Figure()
    
    # 添加PE柱状图
    fig.add_trace(go.Bar(
        x=df[company_col],
        y=df[pe_col],
        name=pe_col,
        marker_color='blue'
    ))
    
    # 添加PB柱状图
    fig.add_trace(go.Bar(
        x=df[company_col],
        y=df[pb_col],
        name=pb_col,
        marker_color='red'
    ))
    
    # 设置图表布局
    fig.update_layout(
        title=title,
        xaxis_title="公司",
        yaxis_title="估值倍数",
        height=height,
        barmode='group',
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    return fig

def create_valuation_history_chart(
    df: pd.DataFrame, 
    date_col: str,
    value_col: str,
    percentile_cols: Optional[List[str]] = None,
    title: str = "历史估值分位数", 
    height: int = 500
) -> go.Figure:
    """
    创建历史估值分位数图表
    
    参数:
        df: 包含估值历史数据的DataFrame，必须包含 date_col, value_col 列
        date_col: 日期列名
        value_col: 估值指标列名
        percentile_cols: 分位数列名列表，如 ['10%分位', '50%分位', '90%分位']
        title: 图表标题
        height: 图表高度
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = [date_col, value_col]
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 创建图表
    fig = go.Figure()
    
    # 添加估值指标线
    fig.add_trace(go.Scatter(
        x=df[date_col],
        y=df[value_col],
        mode='lines',
        name=value_col,
        line=dict(color='blue')
    ))
    
    # 添加分位数线
    if percentile_cols:
        colors = ['red', 'green', 'purple', 'orange', 'cyan']
        
        for i, col in enumerate(percentile_cols):
            if col in df.columns:
                color = colors[i % len(colors)]
                
                fig.add_trace(go.Scatter(
                    x=df[date_col],
                    y=df[col],
                    mode='lines',
                    name=col,
                    line=dict(color=color, dash='dash')
                ))
    
    # 设置图表布局
    fig.update_layout(
        title=title,
        xaxis_title="日期",
        yaxis_title=value_col,
        height=height,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    return fig